#include <CL/cl.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>

std::string loadKernelSource(const std::string& filename) {
    std::string kernel_path = KERNELS_DIR "/" + filename;
    std::ifstream file(kernel_path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open kernel file: " + kernel_path);
    }
    return std::string((std::istreambuf_iterator<char>(file)), 
                      std::istreambuf_iterator<char>());
}

int main() {
    cl_int err;

    // --- 获取平台 ---
    cl_uint num_platforms = 0;
    clGetPlatformIDs(0, nullptr, &num_platforms);
    if (num_platforms == 0) {
        std::cerr << "❌ No OpenCL platform found.\n";
        return 1;
    }

    std::vector<cl_platform_id> platforms(num_platforms);
    clGetPlatformIDs(num_platforms, platforms.data(), nullptr);

    cl_platform_id platform = platforms[0];
    std::cout << "✅ Found " << num_platforms << " OpenCL platform(s).\n";

    // --- 获取设备 ---
    cl_uint num_devices = 0;
    clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 0, nullptr, &num_devices);
    if (num_devices == 0) {
        std::cerr << "❌ No GPU device found.\n";
        return 1;
    }

    std::vector<cl_device_id> devices(num_devices);
    clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, num_devices, devices.data(), nullptr);
    cl_device_id device = devices[0];

    char device_name[256];
    clGetDeviceInfo(device, CL_DEVICE_NAME, sizeof(device_name), device_name, nullptr);
    std::cout << "Using device: " << device_name << "\n";

    // --- 检查 cl_khr_subgroups 扩展 ---
    size_t ext_size = 0;
    clGetDeviceInfo(device, CL_DEVICE_EXTENSIONS, 0, nullptr, &ext_size);
    std::string extensions(ext_size, '\0');
    clGetDeviceInfo(device, CL_DEVICE_EXTENSIONS, ext_size, &extensions[0], nullptr);

    if (extensions.find("cl_khr_subgroups") == std::string::npos) {
        std::cerr << "❌ Device does NOT report cl_khr_subgroups.\n";
        return 1;
    }
    std::cout << "✅ Device reports cl_khr_subgroups extension.\n";

    // --- 创建上下文与命令队列 ---
    cl_context context = clCreateContext(nullptr, 1, &device, nullptr, nullptr, &err);
    cl_command_queue queue = clCreateCommandQueueWithProperties(context, device, 0, &err);

    // --- 读取 kernel 源码 ---
    std::string source = loadKernelSource("subgroup.cl");
    const char *src = source.c_str();
    size_t length = source.size();

    // --- 创建并编译程序 ---
    cl_program program = clCreateProgramWithSource(context, 1, &src, &length, &err);
    std::string build_options = "-cl-std=CL2.0";
    err = clBuildProgram(program, 1, &device, build_options.c_str(), nullptr, nullptr);

    if (err != CL_SUCCESS) {
        size_t log_size = 0;
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr, &log_size);
        std::string log(log_size, '\0');
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, log_size, &log[0], nullptr);
        std::cerr << "❌ Kernel build failed:\n" << log << "\n";
        return 1;
    }

    // --- 创建 kernel ---
    cl_kernel kernel = clCreateKernel(program, "subgroup_test", &err);

    // --- 分配缓冲区 ---
    const size_t N = 64;
    cl_mem out_buf = clCreateBuffer(context, CL_MEM_WRITE_ONLY, sizeof(cl_int) * N, nullptr, &err);
    err = clSetKernelArg(kernel, 0, sizeof(cl_mem), &out_buf);

    // --- 启动 kernel ---
    size_t global_size = N;
    size_t local_size = 32;  // 典型工作组大小
    err = clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, &global_size, &local_size, 0, nullptr, nullptr);
    if (err != CL_SUCCESS) {
        std::cerr << "❌ Kernel launch failed (err=" << err << ")\n";
        return 1;
    }

    // --- 读取结果 ---
    std::vector<int> output(N);
    clEnqueueReadBuffer(queue, out_buf, CL_TRUE, 0, sizeof(int) * N, output.data(), 0, nullptr, nullptr);

    std::cout << "✅ Kernel executed successfully. Sample output:\n";
    for (int i = 0; i < 8 && i < N; i++) {
        printf("  out[%2d] = 0x%08x\n", i, output[i]);
    }

    // --- 清理 ---
    clReleaseMemObject(out_buf);
    clReleaseKernel(kernel);
    clReleaseProgram(program);
    clReleaseCommandQueue(queue);
    clReleaseContext(context);

    return 0;
}